{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Persistence\n",
    "Save and load pre-trained models\n",
    "\n",
    "This example demonstrates how to save trained deep-river models to disk and load them later for continued training or inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:53.467791Z",
     "iopub.status.busy": "2025-09-28T08:15:53.467411Z",
     "iopub.status.idle": "2025-09-28T08:15:56.226102Z",
     "shell.execute_reply": "2025-09-28T08:15:56.224780Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from river import datasets, metrics\n",
    "import tempfile\n",
    "import os\n",
    "from deep_river.classification import Classifier\n",
    "from deep_river import base"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create and Train a Classification Model\n",
    "\n",
    "We will create a simple neural network for phishing detection and train it on some data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:56.229314Z",
     "iopub.status.busy": "2025-09-28T08:15:56.228983Z",
     "iopub.status.idle": "2025-09-28T08:15:56.925774Z",
     "shell.execute_reply": "2025-09-28T08:15:56.925098Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><details class=\"river-component river-estimator\"><summary class=\"river-summary\"><pre class=\"river-estimator-name\">Classifier</pre></summary><code class=\"river-estimator-params\">Classifier (\n",
       "  module=SimpleNet(\n",
       "  (dense0): Linear(in_features=9, out_features=5, bias=True)\n",
       "  (nonlin): ReLU()\n",
       "  (dense1): Linear(in_features=5, out_features=2, bias=True)\n",
       "  (softmax): Softmax(dim=-1)\n",
       ")\n",
       "  loss_fn=\"binary_cross_entropy\"\n",
       "  optimizer_fn=\"adam\"\n",
       "  lr=0.001\n",
       "  output_is_logit=True\n",
       "  is_class_incremental=False\n",
       "  is_feature_incremental=False\n",
       "  device=\"cpu\"\n",
       "  seed=42\n",
       ")\n",
       "</code></details><style scoped>\n",
       ".river-estimator {\n",
       "    padding: 1em;\n",
       "    border-style: solid;\n",
       "    background: white;\n",
       "    max-width: max-content;\n",
       "}\n",
       "\n",
       ".river-pipeline {\n",
       "    display: flex;\n",
       "    flex-direction: column;\n",
       "    align-items: center;\n",
       "    background: linear-gradient(#000, #000) no-repeat center / 1.5px 100%;\n",
       "}\n",
       "\n",
       ".river-union {\n",
       "    display: flex;\n",
       "    flex-direction: row;\n",
       "    align-items: center;\n",
       "    justify-content: center;\n",
       "    padding: 1em;\n",
       "    border-style: solid;\n",
       "    background: white;\n",
       "}\n",
       "\n",
       ".river-wrapper {\n",
       "    display: flex;\n",
       "    flex-direction: column;\n",
       "    align-items: center;\n",
       "    justify-content: center;\n",
       "    padding: 1em;\n",
       "    border-style: solid;\n",
       "    background: white;\n",
       "}\n",
       "\n",
       ".river-wrapper > .river-estimator {\n",
       "    margin-top: 1em;\n",
       "}\n",
       "\n",
       "/* Vertical spacing between steps */\n",
       "\n",
       ".river-component + .river-component {\n",
       "    margin-top: 2em;\n",
       "}\n",
       "\n",
       ".river-union > .river-estimator {\n",
       "    margin-top: 0;\n",
       "}\n",
       "\n",
       ".river-union > .river-component {\n",
       "    margin-top: 0;\n",
       "}\n",
       "\n",
       ".river-union > .pipeline {\n",
       "    margin-top: 0;\n",
       "}\n",
       "\n",
       "/* Spacing within a union of estimators */\n",
       "\n",
       ".river-union > .river-component + .river-component {\n",
       "    margin-left: 1em;\n",
       "}\n",
       "\n",
       "/* Typography */\n",
       "\n",
       ".river-estimator-params {\n",
       "    display: block;\n",
       "    white-space: pre-wrap;\n",
       "    font-size: 110%;\n",
       "    margin-top: 1em;\n",
       "}\n",
       "\n",
       ".river-estimator > .river-estimator-params,\n",
       ".river-wrapper > .river-details > river-estimator-params {\n",
       "    background-color: white !important;\n",
       "}\n",
       "\n",
       ".river-wrapper > .river-details {\n",
       "    margin-bottom: 1em;\n",
       "}\n",
       "\n",
       ".river-estimator-name {\n",
       "    display: inline;\n",
       "    margin: 0;\n",
       "    font-size: 110%;\n",
       "}\n",
       "\n",
       "/* Toggle */\n",
       "\n",
       ".river-summary {\n",
       "    display: flex;\n",
       "    align-items:center;\n",
       "    cursor: pointer;\n",
       "}\n",
       "\n",
       ".river-summary > div {\n",
       "    width: 100%;\n",
       "}\n",
       "</style></div>"
      ],
      "text/plain": [
       "Classifier (\n",
       "  module=SimpleNet(\n",
       "  (dense0): Linear(in_features=9, out_features=5, bias=True)\n",
       "  (nonlin): ReLU()\n",
       "  (dense1): Linear(in_features=5, out_features=2, bias=True)\n",
       "  (softmax): Softmax(dim=-1)\n",
       ")\n",
       "  loss_fn=\"binary_cross_entropy\"\n",
       "  optimizer_fn=\"adam\"\n",
       "  lr=0.001\n",
       "  output_is_logit=True\n",
       "  is_class_incremental=False\n",
       "  is_feature_incremental=False\n",
       "  device=\"cpu\"\n",
       "  seed=42\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = datasets.Phishing()\n",
    "metric = metrics.Accuracy()\n",
    "\n",
    "class SimpleNet(nn.Module):\n",
    "    def __init__(self, n_features):\n",
    "        super().__init__()\n",
    "        self.dense0 = nn.Linear(n_features, 5)\n",
    "        self.nonlin = nn.ReLU()\n",
    "        self.dense1 = nn.Linear(5, 2)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, X, **kwargs):\n",
    "        X = self.nonlin(self.dense0(X))\n",
    "        X = self.nonlin(self.dense1(X))\n",
    "        X = self.softmax(X)\n",
    "        return X\n",
    "\n",
    "model = Classifier(\n",
    "    module=SimpleNet(n_features=9),\n",
    "    loss_fn=\"binary_cross_entropy\",\n",
    "    optimizer_fn=\"adam\"\n",
    ")\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the model on 1000 samples from the phishing dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:56.928739Z",
     "iopub.status.busy": "2025-09-28T08:15:56.928422Z",
     "iopub.status.idle": "2025-09-28T08:15:57.870593Z",
     "shell.execute_reply": "2025-09-28T08:15:57.870087Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.725\n"
     ]
    }
   ],
   "source": [
    "# Train the model\n",
    "for x, y in dataset.take(1000):\n",
    "    y_pred = model.predict_one(x)\n",
    "    if y_pred is not None:\n",
    "        metric.update(y, y_pred)\n",
    "    model.learn_one(x, y)\n",
    "    \n",
    "print(f\"Accuracy: {metric.get():.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the Trained Model\n",
    "\n",
    "Now we will save the trained model to disk using the built-in save functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:57.873298Z",
     "iopub.status.busy": "2025-09-28T08:15:57.873095Z",
     "iopub.status.idle": "2025-09-28T08:15:57.882184Z",
     "shell.execute_reply": "2025-09-28T08:15:57.881526Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model saved to: /var/folders/z7/fzv5wnys4hl1pny_xr6c0g5c0000gn/T/tmpp2azp9sr/classifier.pkl\n"
     ]
    }
   ],
   "source": [
    "# Save the model\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "model_path = os.path.join(temp_dir, \"classifier.pkl\")\n",
    "\n",
    "model.save(model_path)\n",
    "print(f\"Model saved to: {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Model and Test\n",
    "\n",
    "Load the saved model and verify it works correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:57.885246Z",
     "iopub.status.busy": "2025-09-28T08:15:57.884896Z",
     "iopub.status.idle": "2025-09-28T08:15:57.894575Z",
     "shell.execute_reply": "2025-09-28T08:15:57.893618Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model loaded: Classifier\n",
      "Prediction: True, Expected: True\n"
     ]
    }
   ],
   "source": [
    "# Load the model\n",
    "loaded_model = Classifier.load(model_path)\n",
    "print(f\"Model loaded: {type(loaded_model).__name__}\")\n",
    "\n",
    "\n",
    "x, y = next(dataset.take(1))\n",
    "y_pred = loaded_model.predict_one(x)\n",
    "\n",
    "print(f\"Prediction: {y_pred}, Expected: {y}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "Clean up temporary files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-28T08:15:57.897122Z",
     "iopub.status.busy": "2025-09-28T08:15:57.896920Z",
     "iopub.status.idle": "2025-09-28T08:15:57.900525Z",
     "shell.execute_reply": "2025-09-28T08:15:57.899832Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cleanup completed\n"
     ]
    }
   ],
   "source": [
    "# Clean up\n",
    "import shutil\n",
    "shutil.rmtree(temp_dir)\n",
    "print(\"Cleanup completed\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deep-river",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
